package com.iking.toolkit.util;

import com.benjaminwan.ocrlibrary.OcrResult;
import com.iking.toolkit.model.Msg;
import com.iking.toolkit.util.ocr.DrawImageAreaModel;
import com.iking.toolkit.util.ocr.OcrRecognizeModel;
import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.List;

/**
 * 功能描述: ocr识别工具类
 *
 * @author: yfl
 * @date: 2024/7/17 16:16
 */
@Slf4j
public class OCRUtils {
    private static InferenceEngine engine = InferenceEngine.getInstance(Model.ONNX_PPOCR_V4);

    /**
     * 对图片进行绘制区域操作
     *
     * @param inputImagePath     输入图片的路径
     * @param drawImageAreaModel 包含绘制区域信息的模型
     * @return 返回一个消息对象，包含操作结果信息
     */
    public static Msg<BufferedImage> drawImageArea(String inputImagePath, DrawImageAreaModel drawImageAreaModel) {
        try {
            // 加载图片
            BufferedImage image = ImageIO.read(new File(inputImagePath));
            return Msg.sucess(createImageArea(image, drawImageAreaModel));
        } catch (IOException e) {
            log.error("绘制图片区域异常: 异常信息:", e);
            return Msg.fail("绘制图片区域异常!");
        }
    }

    /**
     * 对图片进行绘制区域操作
     *
     * @param file               输入图片的对象
     * @param drawImageAreaModel 包含绘制区域信息的模型
     * @return 返回一个消息对象，包含操作结果信息
     */
    public static Msg<BufferedImage> drawImageArea(File file, DrawImageAreaModel drawImageAreaModel) {
        try {
            // 加载图片
            BufferedImage image = ImageIO.read(file);
            return Msg.sucess(createImageArea(image, drawImageAreaModel));
        } catch (IOException e) {
            log.error("绘制图片区域异常: 异常信息:", e);
            return Msg.fail("绘制图片区域异常!");
        }
    }

    /**
     * 对图片进行绘制区域操作
     *
     * @param bufferedImage      输入图片bufferedImage
     * @param drawImageAreaModel 包含绘制区域信息的模型
     * @return 返回一个消息对象，包含操作结果信息
     */
    public static Msg<BufferedImage> drawImageArea(BufferedImage bufferedImage, DrawImageAreaModel drawImageAreaModel) {
        // 加载图片
        return Msg.sucess(createImageArea(bufferedImage, drawImageAreaModel));
    }

    /**
     * 对图片进行绘制区域操作
     *
     * @param inputStream        输入图片inputStream流
     * @param drawImageAreaModel 包含绘制区域信息的模型
     * @return 返回一个消息对象，包含操作结果信息
     */
    public static Msg<BufferedImage> drawImageArea(InputStream inputStream, DrawImageAreaModel drawImageAreaModel) {
        try {
            // 加载图片
            BufferedImage image = ImageIO.read(inputStream);
            // 加载图片
            return Msg.sucess(createImageArea(image, drawImageAreaModel));
        } catch (IOException e) {
            log.error("绘制图片区域异常: 异常信息:", e);
            return Msg.fail("绘制图片区域异常!");
        }
    }

    /**
     * 功能描述: 绘制图片区域
     *
     * @param: [bufferedImage, drawImageAreaModel]
     * @author: yfl
     * @date: 2024/8/8 15:57
     */
    private static BufferedImage createImageArea(BufferedImage image, DrawImageAreaModel drawImageAreaModel) {
        // 创建 Graphics2D 对象用于绘制
        Graphics2D g2d = image.createGraphics();
        g2d.setColor(drawImageAreaModel.getColor());
        if (drawImageAreaModel.getRectangles() != null) {
            for (Rectangle rectangle : drawImageAreaModel.getRectangles()) {
                g2d.fillRect((int) rectangle.getX(), (int) rectangle.getY(), (int) rectangle.getWidth(), (int) rectangle.getHeight());
            }
        }
        // 释放资源
        g2d.dispose();
        return image;
    }

    /**
     * 功能描述: ocr识别图片文字(上传文件)
     *
     * @param: [multipartFile]
     * @author: yfl
     * @date: 2024/7/17 17:23
     */
    public static Msg<OcrResult> ocrRecognizeImage(MultipartFile multipartFile) {
        File file = null;
        try {
            file = FileFormatUtil.multipartFileToFile(multipartFile);
            OcrResult ocrResult = engine.runOcr(file.getPath());
            return Msg.sucess(ocrResult);
        } catch (Exception e) {
            file.deleteOnExit();
            log.error("ocr识别异常: 异常信息:", e);
            return Msg.fail("ocr识别异常!");
        }finally {
            // 关键改进：在finally块中确保文件被删除，无论成功失败都会执行
            if (file != null && file.exists()) {
                // 尝试删除文件
                boolean deleted = file.delete();
                if (!deleted) {
                    // 若删除失败，记录警告日志
                    log.warn("临时文件删除失败，路径：{}", file.getAbsolutePath());
                    // 作为备选方案，设置延迟删除（防止当前有进程占用文件）
                    file.deleteOnExit();
                }
            }
        }
    }

    /**
     * 功能描述: ocr识别图片文字(文件流)
     *
     * @param: [file]
     * @author: yfl
     * @date: 2024/7/17 17:23
     */
    public static Msg<OcrResult> ocrRecognizeImage(File file) {
        try {
            OcrResult ocrResult = engine.runOcr(file.getPath());
            return Msg.sucess(ocrResult);
        } catch (Exception e) {
            log.error("ocr识别异常: 异常信息:", e);
            return Msg.fail("ocr识别异常!");
        }
    }

    /**
     * 功能描述: ocr识别图片文字(本地地址)
     *
     * @param: [filePath]
     * @author: yfl
     * @date: 2024/7/17 17:23
     */
    public static Msg<OcrResult> ocrRecognizeImage(String filePath) {
        try {
            OcrResult ocrResult = engine.runOcr(filePath);
            return Msg.sucess(ocrResult);
        } catch (Exception e) {
            log.error("ocr识别异常: 异常信息:", e);
            return Msg.fail("ocr识别异常!");
        }
    }

    /**
     * 识别图片区域中的文字(上传文件)
     * 该方法主要用于图片的OCR识别，可以从图片文件中识别出指定区域的文字内容
     *
     * @param multipartFile         图片文件File用于识别
     * @param ocrRecognizeModelList 一个Rectangle的列表，指定图片中需要识别的区域
     * @return 返回一个Msg对象，包含识别结果列表，每个识别结果对应一个识别出的区域
     * <p>
     * 注意：该方法不处理图片的加载和显示，只负责对图片进行文字识别
     */
    public static Msg<List<OcrRecognizeModel>> ocrRecognizeImage(MultipartFile multipartFile, List<OcrRecognizeModel> ocrRecognizeModelList) {
        File file = null;
        try {
            file = FileFormatUtil.multipartFileToFile(multipartFile);
            BufferedImage bufferedImage = ImageIO.read(file);
            if (ocrRecognizeModelList == null || ocrRecognizeModelList.size() < 1) {
                return Msg.fail("请传入需要识别的区域!");
            }
            return Msg.sucess(ocrRecognizeModelList(bufferedImage, ocrRecognizeModelList));
        } catch (IOException e) {
            file.deleteOnExit();
            log.error("识别图片区域中的文字异常: 异常信息=", e);
            return Msg.fail("识别图片区域中的文字异常!");
        }finally {
            // 关键改进：在finally块中确保文件被删除，无论成功失败都会执行
            if (file != null && file.exists()) {
                // 尝试删除文件
                boolean deleted = file.delete();
                if (!deleted) {
                    // 若删除失败，记录警告日志
                    log.warn("临时文件删除失败，路径：{}", file.getAbsolutePath());
                    // 作为备选方案，设置延迟删除（防止当前有进程占用文件）
                    file.deleteOnExit();
                }
            }
        }
    }

    /**
     * 识别图片区域中的文字(文件流)
     * 该方法主要用于图片的OCR识别，可以从图片文件中识别出指定区域的文字内容
     *
     * @param file                  图片文件File用于识别
     * @param ocrRecognizeModelList 一个Rectangle的列表，指定图片中需要识别的区域
     * @return 返回一个Msg对象，包含识别结果列表，每个识别结果对应一个识别出的区域
     * <p>
     * 注意：该方法不处理图片的加载和显示，只负责对图片进行文字识别
     */
    public static Msg<List<OcrRecognizeModel>> ocrRecognizeImage(File file, List<OcrRecognizeModel> ocrRecognizeModelList) {
        try {
            BufferedImage bufferedImage = ImageIO.read(file);
            if (ocrRecognizeModelList == null || ocrRecognizeModelList.size() < 1) {
                return Msg.fail("请传入需要识别的区域!");
            }
            return Msg.sucess(ocrRecognizeModelList(bufferedImage, ocrRecognizeModelList));
        } catch (IOException e) {
            log.error("识别图片区域中的文字异常: 异常信息=", e);
            return Msg.fail("识别图片区域中的文字异常!");
        }
    }

    /**
     * 识别图片区域中的文字(本地地址)
     * 该方法主要用于图片的OCR识别，可以从图片文件中识别出指定区域的文字内容
     *
     * @param filePath              图片文件的路径用于识别
     * @param ocrRecognizeModelList 一个Rectangle的列表，指定图片中需要识别的区域
     * @return 返回一个Msg对象，包含识别结果列表，每个识别结果对应一个识别出的区域
     * <p>
     * 注意：该方法不处理图片的加载和显示，只负责对图片进行文字识别
     */
    public static Msg<List<OcrRecognizeModel>> ocrRecognizeImage(String filePath, List<OcrRecognizeModel> ocrRecognizeModelList) {
        try {
            BufferedImage bufferedImage = ImageIO.read(new File(filePath));
            if (ocrRecognizeModelList == null || ocrRecognizeModelList.size() < 1) {
                return Msg.fail("请传入需要识别的区域!");
            }
            return Msg.sucess(ocrRecognizeModelList(bufferedImage, ocrRecognizeModelList));
        } catch (IOException e) {
            log.error("识别图片区域中的文字异常: 异常信息=", e);
            return Msg.fail("识别图片区域中的文字异常!");
        }
    }

    /**
     * 功能描述: 对图片进行区域Ocr
     *
     * @param: [bufferedImage, ocrRecognizeModelList]
     * @author: yfl
     * @date: 2024/8/7 15:39
     */
    private static List<OcrRecognizeModel> ocrRecognizeModelList(BufferedImage bufferedImage, List<OcrRecognizeModel> ocrRecognizeModelList) throws IOException {
        for (OcrRecognizeModel ocrRecognizeModel : ocrRecognizeModelList) {
            BufferedImage croppedImage = bufferedImage.getSubimage(ocrRecognizeModel.getOcrAreaModel().getX(), ocrRecognizeModel.getOcrAreaModel().getY(), ocrRecognizeModel.getOcrAreaModel().getWidth(), ocrRecognizeModel.getOcrAreaModel().getHeight());
            OcrResult ocrResult = OCRUtils.ocrRecognizeImage(FileFormatUtil.bufferedImageToFile(croppedImage, "png")).getObj();
            ocrRecognizeModel.setOcrResult(ocrResult);
        }
        return ocrRecognizeModelList;
    }

    public static void main(String[] args) throws IOException {
        Color color = new Color(255, 0, 0, 125 ); // 红色半透明
        BufferedImage image = ImageIO.read(new File("E://PDF/d.png"));
        Msg<BufferedImage> drawImageArea = OCRUtils.drawImageArea(image, new DrawImageAreaModel(color, Arrays.asList(new Rectangle(300, 300, 300, 100))));
        // 保存修改后的图片
        System.out.println(drawImageArea.getObj().getType());
        ImageIO.write(drawImageArea.getObj(), "png", new File("E://PDF/d-area.png"));
//         List<OcrRecognizeModel> ocrRecognizeModelList = new ArrayList<>();
//         ocrRecognizeModelList.add(new OcrRecognizeModel("A1", new Rectangle(300, 300, 300, 100)));
// //        ocrRecognizeModelList.add(new OcrRecognizeModel("A2", new Rectangle(50, 150, 200, 500)));
//         Msg<List<OcrRecognizeModel>> ocrRecognizeImage = ocrRecognizeImage("E://PDF/image1.png", ocrRecognizeModelList);
//         System.out.println("结果:" + JSONUtil.parse(ocrRecognizeImage.getObj()));

    }
}
